{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3ca73751-8c0c-4bd5-88fe-78b856103cca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "import subprocess\n",
    "import sys\n",
    "import time\n",
    "from copy import deepcopy\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.distributed as dist\n",
    "import torch.nn as nn\n",
    "import yaml\n",
    "from torch.optim import lr_scheduler\n",
    "from tqdm import tqdm\n",
    "from models.experimental import attempt_load\n",
    "from models.yolo import Model\n",
    "from utils.autoanchor import check_anchors\n",
    "from utils.autobatch import check_train_batch_size\n",
    "from utils.callbacks import Callbacks\n",
    "from utils.dataloaders import create_dataloader\n",
    "from utils.downloads import attempt_download, is_url\n",
    "from utils.general import (\n",
    "    LOGGER,\n",
    "    TQDM_BAR_FORMAT,\n",
    "    check_amp,\n",
    "    check_dataset,\n",
    "    check_file,\n",
    "    check_git_info,\n",
    "    check_git_status,\n",
    "    check_img_size,\n",
    "    check_requirements,\n",
    "    check_suffix,\n",
    "    check_yaml,\n",
    "    colorstr,\n",
    "    get_latest_run,\n",
    "    increment_path,\n",
    "    init_seeds,\n",
    "    intersect_dicts,\n",
    "    labels_to_class_weights,\n",
    "    labels_to_image_weights,\n",
    "    methods,\n",
    "    one_cycle,\n",
    "    print_args,\n",
    "    print_mutation,\n",
    "    strip_optimizer,\n",
    "    yaml_save,\n",
    ")\n",
    "from utils.loggers import LOGGERS, Loggers\n",
    "from utils.loggers.comet.comet_utils import check_comet_resume\n",
    "from utils.loss import ComputeLoss\n",
    "from utils.metrics import fitness\n",
    "from utils.plots import plot_evolve\n",
    "from utils.torch_utils import (\n",
    "    EarlyStopping,\n",
    "    ModelEMA,\n",
    "    de_parallel,\n",
    "    select_device,\n",
    "    smart_DDP,\n",
    "    smart_optimizer,\n",
    "    smart_resume,\n",
    "    torch_distributed_zero_first,\n",
    ")\n",
    "LOCAL_RANK = int(os.getenv(\"LOCAL_RANK\", -1))  # https://pytorch.org/docs/stable/elastic/run.html\n",
    "RANK = int(os.getenv(\"RANK\", -1))\n",
    "WORLD_SIZE = int(os.getenv(\"WORLD_SIZE\", 1))\n",
    "GIT_INFO = check_git_info()\n",
    "from pathlib import Path\n",
    "\n",
    "import torch.optim as optim\n",
    "from EWC_module.EWC import compute_ewc_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2e03f945-ffbd-421a-a45a-38b0bf1e92ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "                 from  n    params  module                                  arguments                     \n",
      "  0                -1  1      3520  models.common.Conv                      [3, 32, 6, 2, 2]              \n",
      "  1                -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                \n",
      "  2                -1  1     18816  models.common.C3                        [64, 64, 1]                   \n",
      "  3                -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               \n",
      "  4                -1  2    115712  models.common.C3                        [128, 128, 2]                 \n",
      "  5                -1  1    295424  models.common.Conv                      [128, 256, 3, 2]              \n",
      "  6                -1  3    625152  models.common.C3                        [256, 256, 3]                 \n",
      "  7                -1  1   1180672  models.common.Conv                      [256, 512, 3, 2]              \n",
      "  8                -1  1   1182720  models.common.C3                        [512, 512, 1]                 \n",
      "  9                -1  1    656896  models.common.SPPF                      [512, 512, 5]                 \n",
      " 10                -1  1    131584  models.common.Conv                      [512, 256, 1, 1]              \n",
      " 11                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          \n",
      " 12           [-1, 6]  1         0  models.common.Concat                    [1]                           \n",
      " 13                -1  1    361984  models.common.C3                        [512, 256, 1, False]          \n",
      " 14                -1  1     33024  models.common.Conv                      [256, 128, 1, 1]              \n",
      " 15                -1  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          \n",
      " 16           [-1, 4]  1         0  models.common.Concat                    [1]                           \n",
      " 17                -1  1     90880  models.common.C3                        [256, 128, 1, False]          \n",
      " 18                -1  1    147712  models.common.Conv                      [128, 128, 3, 2]              \n",
      " 19          [-1, 14]  1         0  models.common.Concat                    [1]                           \n",
      " 20                -1  1    296448  models.common.C3                        [256, 256, 1, False]          \n",
      " 21                -1  1    590336  models.common.Conv                      [256, 256, 3, 2]              \n",
      " 22          [-1, 10]  1         0  models.common.Concat                    [1]                           \n",
      " 23                -1  1   1182720  models.common.C3                        [512, 512, 1, False]          \n",
      " 24      [17, 20, 23]  1     35061  models.yolo.Detect                      [8, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]\n",
      "YOLOv5s_kitti summary: 214 layers, 7041205 parameters, 7041205 gradients, 16.0 GFLOPs\n",
      "\n"
     ]
    }
   ],
   "source": [
    "ckpt = torch.load('runs/train/exp3/weights/best.pt', map_location=\"cpu\")  # load checkpoint to CPU to avoid CUDA memory leak\n",
    "opt = ckpt['opt']\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "save_dir, data, hyp, resume, cfg, single_cls = (\n",
    "    opt['save_dir'],\n",
    "    opt['data'],\n",
    "    opt['hyp'],\n",
    "    opt['resume'],\n",
    "    opt['cfg'],\n",
    "    opt['single_cls'],\n",
    ")\n",
    "cfg = 'models/yolov5s_kitti.yaml'\n",
    "model = Model(cfg, ch=3, nc=8, anchors=hyp.get(\"anchors\")).to(device)\n",
    "model.nc = 8  # attach number of classes to model\n",
    "model.hyp = hyp  # attach hyperparameters to model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14322750-7f2e-47ba-a396-8da118d1f111",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    print(param)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7beadce5-a590-4945-a3ac-a69e71af9d7b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "d94f14ae-3a90-4a59-961f-7d0573a9783a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['model.0.conv.weight', 'model.0.bn.weight', 'model.0.bn.bias', 'model.0.bn.running_mean', 'model.0.bn.running_var', 'model.0.bn.num_batches_tracked', 'model.1.conv.weight', 'model.1.bn.weight', 'model.1.bn.bias', 'model.1.bn.running_mean', 'model.1.bn.running_var', 'model.1.bn.num_batches_tracked', 'model.2.cv1.conv.weight', 'model.2.cv1.bn.weight', 'model.2.cv1.bn.bias', 'model.2.cv1.bn.running_mean', 'model.2.cv1.bn.running_var', 'model.2.cv1.bn.num_batches_tracked', 'model.2.cv2.conv.weight', 'model.2.cv2.bn.weight', 'model.2.cv2.bn.bias', 'model.2.cv2.bn.running_mean', 'model.2.cv2.bn.running_var', 'model.2.cv2.bn.num_batches_tracked', 'model.2.cv3.conv.weight', 'model.2.cv3.bn.weight', 'model.2.cv3.bn.bias', 'model.2.cv3.bn.running_mean', 'model.2.cv3.bn.running_var', 'model.2.cv3.bn.num_batches_tracked', 'model.2.m.0.cv1.conv.weight', 'model.2.m.0.cv1.bn.weight', 'model.2.m.0.cv1.bn.bias', 'model.2.m.0.cv1.bn.running_mean', 'model.2.m.0.cv1.bn.running_var', 'model.2.m.0.cv1.bn.num_batches_tracked', 'model.2.m.0.cv2.conv.weight', 'model.2.m.0.cv2.bn.weight', 'model.2.m.0.cv2.bn.bias', 'model.2.m.0.cv2.bn.running_mean', 'model.2.m.0.cv2.bn.running_var', 'model.2.m.0.cv2.bn.num_batches_tracked', 'model.3.conv.weight', 'model.3.bn.weight', 'model.3.bn.bias', 'model.3.bn.running_mean', 'model.3.bn.running_var', 'model.3.bn.num_batches_tracked', 'model.4.cv1.conv.weight', 'model.4.cv1.bn.weight', 'model.4.cv1.bn.bias', 'model.4.cv1.bn.running_mean', 'model.4.cv1.bn.running_var', 'model.4.cv1.bn.num_batches_tracked', 'model.4.cv2.conv.weight', 'model.4.cv2.bn.weight', 'model.4.cv2.bn.bias', 'model.4.cv2.bn.running_mean', 'model.4.cv2.bn.running_var', 'model.4.cv2.bn.num_batches_tracked', 'model.4.cv3.conv.weight', 'model.4.cv3.bn.weight', 'model.4.cv3.bn.bias', 'model.4.cv3.bn.running_mean', 'model.4.cv3.bn.running_var', 'model.4.cv3.bn.num_batches_tracked', 'model.4.m.0.cv1.conv.weight', 'model.4.m.0.cv1.bn.weight', 'model.4.m.0.cv1.bn.bias', 'model.4.m.0.cv1.bn.running_mean', 'model.4.m.0.cv1.bn.running_var', 'model.4.m.0.cv1.bn.num_batches_tracked', 'model.4.m.0.cv2.conv.weight', 'model.4.m.0.cv2.bn.weight', 'model.4.m.0.cv2.bn.bias', 'model.4.m.0.cv2.bn.running_mean', 'model.4.m.0.cv2.bn.running_var', 'model.4.m.0.cv2.bn.num_batches_tracked', 'model.4.m.1.cv1.conv.weight', 'model.4.m.1.cv1.bn.weight', 'model.4.m.1.cv1.bn.bias', 'model.4.m.1.cv1.bn.running_mean', 'model.4.m.1.cv1.bn.running_var', 'model.4.m.1.cv1.bn.num_batches_tracked', 'model.4.m.1.cv2.conv.weight', 'model.4.m.1.cv2.bn.weight', 'model.4.m.1.cv2.bn.bias', 'model.4.m.1.cv2.bn.running_mean', 'model.4.m.1.cv2.bn.running_var', 'model.4.m.1.cv2.bn.num_batches_tracked', 'model.5.conv.weight', 'model.5.bn.weight', 'model.5.bn.bias', 'model.5.bn.running_mean', 'model.5.bn.running_var', 'model.5.bn.num_batches_tracked', 'model.6.cv1.conv.weight', 'model.6.cv1.bn.weight', 'model.6.cv1.bn.bias', 'model.6.cv1.bn.running_mean', 'model.6.cv1.bn.running_var', 'model.6.cv1.bn.num_batches_tracked', 'model.6.cv2.conv.weight', 'model.6.cv2.bn.weight', 'model.6.cv2.bn.bias', 'model.6.cv2.bn.running_mean', 'model.6.cv2.bn.running_var', 'model.6.cv2.bn.num_batches_tracked', 'model.6.cv3.conv.weight', 'model.6.cv3.bn.weight', 'model.6.cv3.bn.bias', 'model.6.cv3.bn.running_mean', 'model.6.cv3.bn.running_var', 'model.6.cv3.bn.num_batches_tracked', 'model.6.m.0.cv1.conv.weight', 'model.6.m.0.cv1.bn.weight', 'model.6.m.0.cv1.bn.bias', 'model.6.m.0.cv1.bn.running_mean', 'model.6.m.0.cv1.bn.running_var', 'model.6.m.0.cv1.bn.num_batches_tracked', 'model.6.m.0.cv2.conv.weight', 'model.6.m.0.cv2.bn.weight', 'model.6.m.0.cv2.bn.bias', 'model.6.m.0.cv2.bn.running_mean', 'model.6.m.0.cv2.bn.running_var', 'model.6.m.0.cv2.bn.num_batches_tracked', 'model.6.m.1.cv1.conv.weight', 'model.6.m.1.cv1.bn.weight', 'model.6.m.1.cv1.bn.bias', 'model.6.m.1.cv1.bn.running_mean', 'model.6.m.1.cv1.bn.running_var', 'model.6.m.1.cv1.bn.num_batches_tracked', 'model.6.m.1.cv2.conv.weight', 'model.6.m.1.cv2.bn.weight', 'model.6.m.1.cv2.bn.bias', 'model.6.m.1.cv2.bn.running_mean', 'model.6.m.1.cv2.bn.running_var', 'model.6.m.1.cv2.bn.num_batches_tracked', 'model.6.m.2.cv1.conv.weight', 'model.6.m.2.cv1.bn.weight', 'model.6.m.2.cv1.bn.bias', 'model.6.m.2.cv1.bn.running_mean', 'model.6.m.2.cv1.bn.running_var', 'model.6.m.2.cv1.bn.num_batches_tracked', 'model.6.m.2.cv2.conv.weight', 'model.6.m.2.cv2.bn.weight', 'model.6.m.2.cv2.bn.bias', 'model.6.m.2.cv2.bn.running_mean', 'model.6.m.2.cv2.bn.running_var', 'model.6.m.2.cv2.bn.num_batches_tracked', 'model.7.conv.weight', 'model.7.bn.weight', 'model.7.bn.bias', 'model.7.bn.running_mean', 'model.7.bn.running_var', 'model.7.bn.num_batches_tracked', 'model.8.cv1.conv.weight', 'model.8.cv1.bn.weight', 'model.8.cv1.bn.bias', 'model.8.cv1.bn.running_mean', 'model.8.cv1.bn.running_var', 'model.8.cv1.bn.num_batches_tracked', 'model.8.cv2.conv.weight', 'model.8.cv2.bn.weight', 'model.8.cv2.bn.bias', 'model.8.cv2.bn.running_mean', 'model.8.cv2.bn.running_var', 'model.8.cv2.bn.num_batches_tracked', 'model.8.cv3.conv.weight', 'model.8.cv3.bn.weight', 'model.8.cv3.bn.bias', 'model.8.cv3.bn.running_mean', 'model.8.cv3.bn.running_var', 'model.8.cv3.bn.num_batches_tracked', 'model.8.m.0.cv1.conv.weight', 'model.8.m.0.cv1.bn.weight', 'model.8.m.0.cv1.bn.bias', 'model.8.m.0.cv1.bn.running_mean', 'model.8.m.0.cv1.bn.running_var', 'model.8.m.0.cv1.bn.num_batches_tracked', 'model.8.m.0.cv2.conv.weight', 'model.8.m.0.cv2.bn.weight', 'model.8.m.0.cv2.bn.bias', 'model.8.m.0.cv2.bn.running_mean', 'model.8.m.0.cv2.bn.running_var', 'model.8.m.0.cv2.bn.num_batches_tracked', 'model.9.cv1.conv.weight', 'model.9.cv1.bn.weight', 'model.9.cv1.bn.bias', 'model.9.cv1.bn.running_mean', 'model.9.cv1.bn.running_var', 'model.9.cv1.bn.num_batches_tracked', 'model.9.cv2.conv.weight', 'model.9.cv2.bn.weight', 'model.9.cv2.bn.bias', 'model.9.cv2.bn.running_mean', 'model.9.cv2.bn.running_var', 'model.9.cv2.bn.num_batches_tracked', 'model.10.conv.weight', 'model.10.bn.weight', 'model.10.bn.bias', 'model.10.bn.running_mean', 'model.10.bn.running_var', 'model.10.bn.num_batches_tracked', 'model.13.cv1.conv.weight', 'model.13.cv1.bn.weight', 'model.13.cv1.bn.bias', 'model.13.cv1.bn.running_mean', 'model.13.cv1.bn.running_var', 'model.13.cv1.bn.num_batches_tracked', 'model.13.cv2.conv.weight', 'model.13.cv2.bn.weight', 'model.13.cv2.bn.bias', 'model.13.cv2.bn.running_mean', 'model.13.cv2.bn.running_var', 'model.13.cv2.bn.num_batches_tracked', 'model.13.cv3.conv.weight', 'model.13.cv3.bn.weight', 'model.13.cv3.bn.bias', 'model.13.cv3.bn.running_mean', 'model.13.cv3.bn.running_var', 'model.13.cv3.bn.num_batches_tracked', 'model.13.m.0.cv1.conv.weight', 'model.13.m.0.cv1.bn.weight', 'model.13.m.0.cv1.bn.bias', 'model.13.m.0.cv1.bn.running_mean', 'model.13.m.0.cv1.bn.running_var', 'model.13.m.0.cv1.bn.num_batches_tracked', 'model.13.m.0.cv2.conv.weight', 'model.13.m.0.cv2.bn.weight', 'model.13.m.0.cv2.bn.bias', 'model.13.m.0.cv2.bn.running_mean', 'model.13.m.0.cv2.bn.running_var', 'model.13.m.0.cv2.bn.num_batches_tracked', 'model.14.conv.weight', 'model.14.bn.weight', 'model.14.bn.bias', 'model.14.bn.running_mean', 'model.14.bn.running_var', 'model.14.bn.num_batches_tracked', 'model.17.cv1.conv.weight', 'model.17.cv1.bn.weight', 'model.17.cv1.bn.bias', 'model.17.cv1.bn.running_mean', 'model.17.cv1.bn.running_var', 'model.17.cv1.bn.num_batches_tracked', 'model.17.cv2.conv.weight', 'model.17.cv2.bn.weight', 'model.17.cv2.bn.bias', 'model.17.cv2.bn.running_mean', 'model.17.cv2.bn.running_var', 'model.17.cv2.bn.num_batches_tracked', 'model.17.cv3.conv.weight', 'model.17.cv3.bn.weight', 'model.17.cv3.bn.bias', 'model.17.cv3.bn.running_mean', 'model.17.cv3.bn.running_var', 'model.17.cv3.bn.num_batches_tracked', 'model.17.m.0.cv1.conv.weight', 'model.17.m.0.cv1.bn.weight', 'model.17.m.0.cv1.bn.bias', 'model.17.m.0.cv1.bn.running_mean', 'model.17.m.0.cv1.bn.running_var', 'model.17.m.0.cv1.bn.num_batches_tracked', 'model.17.m.0.cv2.conv.weight', 'model.17.m.0.cv2.bn.weight', 'model.17.m.0.cv2.bn.bias', 'model.17.m.0.cv2.bn.running_mean', 'model.17.m.0.cv2.bn.running_var', 'model.17.m.0.cv2.bn.num_batches_tracked', 'model.18.conv.weight', 'model.18.bn.weight', 'model.18.bn.bias', 'model.18.bn.running_mean', 'model.18.bn.running_var', 'model.18.bn.num_batches_tracked', 'model.20.cv1.conv.weight', 'model.20.cv1.bn.weight', 'model.20.cv1.bn.bias', 'model.20.cv1.bn.running_mean', 'model.20.cv1.bn.running_var', 'model.20.cv1.bn.num_batches_tracked', 'model.20.cv2.conv.weight', 'model.20.cv2.bn.weight', 'model.20.cv2.bn.bias', 'model.20.cv2.bn.running_mean', 'model.20.cv2.bn.running_var', 'model.20.cv2.bn.num_batches_tracked', 'model.20.cv3.conv.weight', 'model.20.cv3.bn.weight', 'model.20.cv3.bn.bias', 'model.20.cv3.bn.running_mean', 'model.20.cv3.bn.running_var', 'model.20.cv3.bn.num_batches_tracked', 'model.20.m.0.cv1.conv.weight', 'model.20.m.0.cv1.bn.weight', 'model.20.m.0.cv1.bn.bias', 'model.20.m.0.cv1.bn.running_mean', 'model.20.m.0.cv1.bn.running_var', 'model.20.m.0.cv1.bn.num_batches_tracked', 'model.20.m.0.cv2.conv.weight', 'model.20.m.0.cv2.bn.weight', 'model.20.m.0.cv2.bn.bias', 'model.20.m.0.cv2.bn.running_mean', 'model.20.m.0.cv2.bn.running_var', 'model.20.m.0.cv2.bn.num_batches_tracked', 'model.21.conv.weight', 'model.21.bn.weight', 'model.21.bn.bias', 'model.21.bn.running_mean', 'model.21.bn.running_var', 'model.21.bn.num_batches_tracked', 'model.23.cv1.conv.weight', 'model.23.cv1.bn.weight', 'model.23.cv1.bn.bias', 'model.23.cv1.bn.running_mean', 'model.23.cv1.bn.running_var', 'model.23.cv1.bn.num_batches_tracked', 'model.23.cv2.conv.weight', 'model.23.cv2.bn.weight', 'model.23.cv2.bn.bias', 'model.23.cv2.bn.running_mean', 'model.23.cv2.bn.running_var', 'model.23.cv2.bn.num_batches_tracked', 'model.23.cv3.conv.weight', 'model.23.cv3.bn.weight', 'model.23.cv3.bn.bias', 'model.23.cv3.bn.running_mean', 'model.23.cv3.bn.running_var', 'model.23.cv3.bn.num_batches_tracked', 'model.23.m.0.cv1.conv.weight', 'model.23.m.0.cv1.bn.weight', 'model.23.m.0.cv1.bn.bias', 'model.23.m.0.cv1.bn.running_mean', 'model.23.m.0.cv1.bn.running_var', 'model.23.m.0.cv1.bn.num_batches_tracked', 'model.23.m.0.cv2.conv.weight', 'model.23.m.0.cv2.bn.weight', 'model.23.m.0.cv2.bn.bias', 'model.23.m.0.cv2.bn.running_mean', 'model.23.m.0.cv2.bn.running_var', 'model.23.m.0.cv2.bn.num_batches_tracked', 'model.24.anchors', 'model.24.m.0.weight', 'model.24.m.0.bias', 'model.24.m.1.weight', 'model.24.m.1.bias', 'model.24.m.2.weight', 'model.24.m.2.bias'])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ckpt[\"model\"].float().state_dict().keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d31291a2-1154-4f77-ba3a-d87badf94866",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['epoch', 'best_fitness', 'model', 'ema', 'updates', 'optimizer', 'opt', 'git', 'date'])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ckpt.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0764e8a4-95b5-49b2-b0f8-fd4d7337c7a7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "584e3c41-ea7f-4133-8b43-7308b2a4053d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9f465e9-8f2d-4127-91ae-64425e959d0e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d7969e73-dd10-4f60-a4cf-c97263a2930e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ewc_temp = torch.load('runs/train/exp3/weights/fisher.pt')\n",
    "ewc_pack = {\n",
    "    'ewc_enable' : True,\n",
    "    'fisher_matrix' : ewc_temp['fisher_matrix'],\n",
    "    'old_modelpara' : ewc_temp['old_modelpara'],\n",
    "    'model' : None, \n",
    "    'ewc_lambda': 1e4, \n",
    "    'device':device, \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5e149721-ae8e-45b9-b18a-bff948fad0e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.model.parameters(), 0.001)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0ee99d01-b245-4082-af40-850db83ff2ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EWC loss at batch 0: 59798172.0\n",
      "EWC loss at batch 20: 41545832.0\n",
      "EWC loss at batch 40: 34795740.0\n",
      "EWC loss at batch 60: 31386618.0\n",
      "EWC loss at batch 80: 29145402.0\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(100):  # 训练5个epoch\n",
    "    ewc_pack['model'] = model.model\n",
    "    ewc_loss = compute_ewc_loss(ewc_pack)\n",
    "    if epoch % 20 == 0:\n",
    "        print(f\"EWC loss at batch {epoch}: {ewc_loss.item()}\")\n",
    "    optimizer.zero_grad()\n",
    "    ewc_loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a0def3c-52e5-46c1-a28b-36bebeb00506",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, param in model.model.named_parameters():\n",
    "    if name in ewc_pack['old_modelpara']:  # 检查是否有对应的旧参数\n",
    "        param.data.copy_(ewc_pack['old_modelpara'][name].data)  # 将旧模型参数赋值给当前模型\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45cc610a-11b5-4e71-a373-2f4778397cb5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "c23aba91-6d5c-423a-8d92-76f524c7ae4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt = {\n",
    "    \"epoch\": -1,\n",
    "    \"model\": deepcopy(de_parallel(model)).half(),\n",
    "    \"optimizer\": optimizer.state_dict(),\n",
    "}\n",
    "\n",
    "torch.save(ckpt, 'test.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90004d1b-8395-4307-bda6-5db02b6e39ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = 'test.pt'\n",
    "\n",
    "val_command = f\" \\\n",
    "python val.py \\\n",
    "--data data/kitti.yaml \\\n",
    "--weights {model_path} \\\n",
    "--task test &&\\\n",
    "echo 'Test set val successfully!' \\\n",
    "\" \n",
    "!{val_command}\n",
    "\n",
    "# 这个是没有ewc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c6cb798-69f3-4dab-ba4a-42c1a0c9b919",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d49bede-59a0-436b-b464-f1a6d8f15d1c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64df61e4-21c0-491b-b882-46c4b10d026e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ab05b65-f90f-4d4c-862b-659fb0af2b64",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
